{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ESMBind模型预测金属结合蛋白的结构实操\n",
    "\n",
    "## 相关链接\n",
    "\n",
    "- [biorxiv预印本:[1] “Predict metal-binding proteins and structures through integration of evolutionary-scale and physics-based modeling | bioRxiv.” Accessed: Aug. 26, 2024. [Online]. Available: Predict metal-binding proteins and structures through integration of evolutionary-scale and physics-based modeling\n",
    "](https://www.biorxiv.org/content/biorxiv/early/2024/08/10/2024.08.09.607368.full.pdf)\n",
    "- [代码仓库](https://github.com/Structurebiology-BNL/ESMBind)\n",
    "\n",
    "## 依赖\n",
    "\n",
    "- Python 3.10+\n",
    "- PyTorch 2.0+\n",
    "- Numba 0.60.0\n",
    "- OpenMM 8.1.1\n",
    "- AmberTools 23.3\n",
    "\n",
    "```bash\n",
    "conda install python=3.10 nvidia::cuda-cudart-dev biopython tqdm numba openmm ambertools pytorch torch-scatter biotite=0.41.0\n",
    "pip install fair-esm torch-geometric multi_modal_binding torchmetrics\n",
    "```\n",
    "\n",
    "本笔记可视化另需安装pymol-open-source.\n",
    "\n",
    "## 代码库协议\n",
    "\n",
    "BSD-3条款.\n",
    "\n",
    "## 代码仓库结构\n",
    "\n",
    "使用`git clone https://github.com/Structurebiology-BNL/ESMBind.git`在指定路径创建esmbind的git克隆仓库."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "esm_bind_path='/home/regen/git_develop/ESMBind' # esmbind git仓库所在位置,自行修改"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1esf.cif  \u001b[0m\u001b[01;34m3D_modeling\u001b[0m/  LICENSE  \u001b[01;34mmulti_modal_binding\u001b[0m/  README.md\n"
     ]
    }
   ],
   "source": [
    "%ls $esm_bind_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[01;34mconfigs\u001b[0m/  \u001b[01;34mdatasets\u001b[0m/             get_esm_if_embedding.py  \u001b[01;34mmodel\u001b[0m/        train.py\n",
      "\u001b[01;34mdata\u001b[0m/     get_esm_embedding.py  inference.py             \u001b[01;34m__pycache__\u001b[0m/  utils.py\n"
     ]
    }
   ],
   "source": [
    "%ls $esm_bind_path/multi_modal_binding/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run-3d-modeling.sh  \u001b[0m\u001b[01;34msrc\u001b[0m/\n"
     ]
    }
   ],
   "source": [
    "%ls $esm_bind_path/3D_modeling/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- multi_modal_binding/:\n",
    "  - configs/:配置文件,规定模型的权重,数据维度\n",
    "    - inference.json:使用预训练后的权重,对训练集中的fasta文件进行推断\n",
    "    - training.json:随机初始化模型权重,用训练集中的fasta文件进行训练\n",
    "  - datasets/:FASTA多序列文件和esm2和esm-if模型预先生成的嵌入数据,可用于训练和预测\n",
    "  - data/:数据处理的相关脚本和常量\n",
    "  - model/:预训练的模型权重参数和模型源代码目录\n",
    "  - get_esm_embedding.py:从序列生成esm嵌入的示例脚本\n",
    "  - get_esm_if_embedding.py:从结构生成esm-if嵌入的示例脚本\n",
    "  - inference.py:载入预训练模型参数,并预测示例序列的残基-金属结合概率\n",
    "  - train.py:以示例序列为输入,训练模型\n",
    "- 3D_modeling/:`run-3d-modeling.sh`生成包含金属离子的cif3D结构,src/main.py为主要使用的脚本,其他脚本为定位/聚类/能量最小化相关模块"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 用法\n",
    "\n",
    "本笔记重点演示使用预训练模型进行预测的步骤,训练步骤从略.\n",
    "\n",
    "### 1.深度学习模型预测金属结合的残基\n",
    "\n",
    "训练和预测金属结合残基的脚本文件和配置文件示例位于`multi_modal_binding/`目录下,可直接进行尝试:\n",
    "\n",
    "```bash\n",
    "python multi_modal_binding/train.py --config multi_modal_binding/configs/training.json # 训练\n",
    "\n",
    "python multi_modal_binding/inference.py --config multi_modal_binding/configs/inference.json # 推断\n",
    "```\n",
    "\n",
    "要使用其他位置的数据作为输入,需修改配置文件.\n",
    "\n",
    "#### 1.1.下载输入数据\n",
    "\n",
    "以肠毒素A(SEA,pdb:1ESF,A链,测定序号25-257aa,UniProtKB:P0A0L2,总长257aa)和SEE(pdb:5FKA,C链,测定序号25-257aa,UniProtKB:P12993,总长257aa)为例,均有Zn2+结合位点,同时是与MHC-II结合的重要位点.\n",
    "\n",
    "可从uniprot下载序列fasta文件:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100   354  100   354    0     0     44      0  0:00:08  0:00:08 --:--:--    87\n",
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100   354  100   354    0     0     22      0  0:00:16  0:00:15  0:00:01   102\n"
     ]
    }
   ],
   "source": [
    "!curl https://rest.uniprot.org/uniprotkb/P0A0L2.fasta -o sea.fasta\n",
    "!curl https://rest.uniprot.org/uniprotkb/P12993.fasta -o see.fasta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fasta序列ID需要和pdb结构文件名(不含前缀)一致,以便进行对应.本笔记修改fasta文件中的id为pdbid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "def alter_fasta_id(filename:str,new_id:str):\n",
    "    '''修改单序列fasta文件中的fasta ID'''\n",
    "    with open(filename,'r') as fasta:\n",
    "        lines=fasta.readlines()\n",
    "        lines[0]=f'>|{new_id}\\n'\n",
    "    with open(filename,'w') as fasta:\n",
    "        fasta.writelines(lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "alter_fasta_id('sea.fasta','sea')\n",
    "alter_fasta_id('see.fasta','see')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "要一次进行多个序列的预测,需要合并为单个fasta文件:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat sea.fasta see.fasta > sea_and_e.fasta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PDB结构推荐从uniprot的alphafold链接下载,因为ESMBind默认pdb文件中的序列与fasta序列一致.SEA/SEE的3D结构已测定,但并未完整测定."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  171k    0  171k    0     0  14032      0 --:--:--  0:00:12 --:--:-- 38438\n",
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  169k    0  169k    0     0  26014      0 --:--:--  0:00:06 --:--:-- 39387\n"
     ]
    }
   ],
   "source": [
    "!curl https://alphafold.ebi.ac.uk/files/AF-P0A0L2-F1-model_v4.pdb -o sea.pdb\n",
    "!curl https://alphafold.ebi.ac.uk/files/AF-P12993-F1-model_v4.pdb -o see.pdb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2 生成序列和结构的嵌入\n",
    "\n",
    "生成esm2模型的嵌入(论文小组提供的脚本会自动尝试下载相应的模型,需要等待一段时间):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: get_esm_embedding.py [-h] [--batch_size BATCH_SIZE]\n",
      "                            [--device {cpu,cuda}] [--gpu GPU]\n",
      "                            fasta_file output_dir\n",
      "\n",
      "Generate protein embeddings using ESM2 (esm2_t33_650M_UR50D)\n",
      "\n",
      "positional arguments:\n",
      "  fasta_file            Path to the input FASTA file\n",
      "  output_dir            Path to the output directory for embeddings\n",
      "\n",
      "options:\n",
      "  -h, --help            show this help message and exit\n",
      "  --batch_size BATCH_SIZE\n",
      "                        Batch size for processing\n",
      "  --device {cpu,cuda}   Device to use (cpu or cuda)\n",
      "  --gpu GPU             GPU device number if using CUDA\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/get_esm_embedding.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "用法:get_esm_embedding.py [-h] [--batch_size BATCH_SIZE]\n",
    "                            [--device {cpu,cuda}] [--gpu GPU]\n",
    "                            fasta_file output_dir\n",
    "\n",
    "- 位置参数:\n",
    "  - fasta_file:输入Fasta文件的位置\n",
    "  - output_dir:输出嵌入的目录\n",
    "- 选项:\n",
    "  - --batch_size BATCH_SIZE 批次大小,默认16\n",
    "  - --device {cpu,cuda}:计算使用的设备,默认优先使用cuda\n",
    "  - --gpu GPU:如果使用CUDA计算,可指定GPU设备编号,默认0\n",
    "\n",
    "试用结果:KeyError:result[\"mask\"],不确定原因,参考源代码后以下面的代码为示例:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Bio import SeqIO\n",
    "import esm\n",
    "import torch\n",
    "import numpy as np\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_esm_embedding(fastafile:str) -> None:\n",
    "    '''在当前目录下创建esm目录并生成指定fasta文件中序列的esm2嵌入(.npy)'''\n",
    "    os.mkdir('esm')\n",
    "    model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "    model = model.to('cpu')\n",
    "    model.eval()\n",
    "\n",
    "    records = list(SeqIO.parse(fastafile, \"fasta\"))\n",
    "    seq_list, id_list = [], []\n",
    "    for rec in records:\n",
    "        seq_list.append(str(rec.seq))\n",
    "        id_list.append(rec.id.split(\"|\")[1])\n",
    "\n",
    "    data = list(zip(id_list, seq_list))\n",
    "    with torch.no_grad():\n",
    "        batch_labels, _, batch_tokens = batch_converter(data)\n",
    "        result = model(\n",
    "                batch_tokens.to('cpu'),\n",
    "                repr_layers=[len(model.layers)],\n",
    "                return_contacts=False,\n",
    "            )\n",
    "        embedding = (\n",
    "            result[\"representations\"][len(model.layers)].detach().cpu().numpy()\n",
    "        )\n",
    "        batch_lens=(batch_tokens != alphabet.padding_idx).sum(1)\n",
    "        for seq_num in range(len(embedding)):\n",
    "            seq_len = batch_lens[seq_num]\n",
    "            # 移除esm嵌入的头尾token\n",
    "            seq_emd = embedding[seq_num][1 : (seq_len - 1)]\n",
    "            np.save(os.path.join('esm',batch_labels[seq_num]), seq_emd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_esm_embedding('sea_and_e.fasta')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "生成esm-if的嵌入:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: get_esm_if_embedding.py [-h] [--chain_id CHAIN_ID]\n",
      "                               fasta_file output_dir pdb_folder\n",
      "\n",
      "Generate protein embeddings from FASTA and PDB files\n",
      "\n",
      "positional arguments:\n",
      "  fasta_file           Path to the input FASTA file\n",
      "  output_dir           Path to the output directory for embeddings\n",
      "  pdb_folder           Path to the folder containing PDB files\n",
      "\n",
      "options:\n",
      "  -h, --help           show this help message and exit\n",
      "  --chain_id CHAIN_ID  Chain ID to use (default: A)\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/get_esm_if_embedding.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "用法:get_esm_if_embedding.py [-h] [--chain_id CHAIN_ID] fasta_file output_dir pdb_folder\n",
    "\n",
    "- 位置参数:\n",
    "  - fasta_file:fasta文件的路径\n",
    "  - output_dir:输出嵌入的目录\n",
    "  - pdb_folder:pdb文件所在路径\n",
    "- 选项:\n",
    "  - -h:显示本帮助信息\n",
    "  - --chain_id CHAIN_ID:使用的链号,默认A(如果多序列fasta文件对应pdb链号不同,需要用单序列fasta文件作为输入,因此官方代码提供的fasta文件示例需要做额外调整)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n",
    "# 禁用GPU,使用GPU会报错"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/dl/lib/python3.11/site-packages/esm/pretrained.py:215: UserWarning: Regression weights not found, predicting contacts will not produce correct results.\n",
      "  warnings.warn(\n",
      "2024-09-12 22:44:23,461 - INFO - Embedding for sea already exists. Skipping.\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/get_esm_if_embedding.py sea.fasta esm_if ./ --chain_id A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/dl/lib/python3.11/site-packages/esm/pretrained.py:215: UserWarning: Regression weights not found, predicting contacts will not produce correct results.\n",
      "  warnings.warn(\n",
      "2024-09-12 22:44:35,933 - INFO - Embedding for see already exists. Skipping.\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/get_esm_if_embedding.py see.fasta esm_if ./ --chain_id A #alphafold预测的结构均为A链,与实验测定编号不同"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3 使用ESMBind预测结合残基\n",
    "\n",
    "由于输入文件位置不同,需要自行生成配置文件:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_config(esm_bind_dir:str,fasta_path:str,esm_embedding_dir:str) -> None:\n",
    "    '''在当前目录生成用于预测的配置文件(inference.json).\n",
    "    \n",
    "    esm_bind_dir:ESMBind仓库所在目录\n",
    "    \n",
    "    fasta_path:多序列文件的路径\n",
    "\n",
    "    esm_embedding_dir:序列和结构嵌入所在目录的父目录(即esm/esm_if目录的上一级)\n",
    "    '''\n",
    "\n",
    "    obj={\n",
    "        \"general\": {\n",
    "            \"seed\": 1,\n",
    "            \"debug\": False\n",
    "        },\n",
    "        \"inference\": {\n",
    "            \"ensemble_path\": f\"{esm_bind_dir}/multi_modal_binding/model/trained_weights\",\n",
    "            \"fasta_path\": fasta_path,\n",
    "            \"precomputed_feature\": esm_embedding_dir\n",
    "        }\n",
    "    }\n",
    "\n",
    "    with open('inference.json','w+') as output:\n",
    "        json.dump(obj,output,indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_config(esm_bind_path,'sea_and_e.fasta','.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: inference.py [-h] [--config CONFIG]\n",
      "\n",
      "options:\n",
      "  -h, --help       show this help message and exit\n",
      "  --config CONFIG  Path of the configuration file\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/inference.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由于源仓库下的`get_esm_embedding.py`与`data_process.py`对fasta id的处理方式不一致,直接调用预测脚本会报错,可使用本仓库下的`data_process.py`替换源仓库下的`multi_modal_binding/data/data_process.py`,默认fasta id格式为\">|序列id\".\n",
    "\n",
    "本问题将在官方仓库进行报告."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 18:21:56,812 | INFO | Output path: multi_modal_binding/results/inference/2024-09-11-18-21\n",
      "2024-09-11 18:21:56,818 | INFO | Test begins at 09-11 18:21\n",
      "/home/regen/git_develop/ESMBind/multi_modal_binding/utils.py:135: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load(ensemble_path + f\"/fold_{i}.pt\", map_location=\"cpu\")\n",
      "2024-09-11 18:21:56,829 | INFO | load weights for fold 1 from /home/regen/git_develop/ESMBind/multi_modal_binding/model/trained_weights\n",
      "2024-09-11 18:21:56,843 | INFO | load weights for fold 2 from /home/regen/git_develop/ESMBind/multi_modal_binding/model/trained_weights\n",
      "2024-09-11 18:21:56,854 | INFO | load weights for fold 3 from /home/regen/git_develop/ESMBind/multi_modal_binding/model/trained_weights\n",
      "2024-09-11 18:21:56,863 | INFO | load weights for fold 4 from /home/regen/git_develop/ESMBind/multi_modal_binding/model/trained_weights\n",
      "2024-09-11 18:21:56,874 | INFO | load weights for fold 5 from /home/regen/git_develop/ESMBind/multi_modal_binding/model/trained_weights\n",
      "2024-09-11 18:21:56,882 | INFO | Total # of test samples for is 2\n",
      "2024-09-11 18:21:59,421 | INFO | Test is done at 09-11 18:21\n",
      "2024-09-11 18:21:59,422 | INFO | Total time used: 2.6\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/multi_modal_binding/inference.py --config inference.json # 如果运行过多次,会自动按运行时间生成子目录,将历史预测放入其中"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在生成的`multi_modal_binding/results/inference/运行时间/`中包含pkl文件就是对应残基结合概率的文件,用于后续的3D建模.\n",
    "\n",
    "也可以使用pickle载入数据,查看哪些残基的概率最高."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from typing import Literal\n",
    "\n",
    "IONS = Literal[\"MG\",\"FE\",\"CU\",\"CO\",\"CA\",\"MN\",\"ZN\"]\n",
    "\n",
    "def best_residues(ion:IONS,pred_file:str,best_n:int=-1) -> dict[str,list[tuple[float,int]]]:\n",
    "    '''对给定预测文件和离子,输出概率最高的前n个残基信息(概率,残基序号).\n",
    "\n",
    "    不指定best_n,或者best_n超过序列长度,将返回全体残基的信息.\n",
    "\n",
    "    编号是序列中的自然序号,从1开始.\n",
    "    '''\n",
    "    output={}\n",
    "    with open(pred_file,'rb') as f:\n",
    "        predictions=pickle.load(f)\n",
    "    if best_n<0:\n",
    "        best_n=max(len(value) for value in predictions['MN'].values())\n",
    "    for fasta_id in predictions[ion]:\n",
    "        pred_with_idx=[(p,idx) for idx,p in enumerate(predictions[ion][fasta_id],1)]\n",
    "        pred_with_idx.sort(reverse=True)\n",
    "        pred_with_idx=pred_with_idx[:best_n]\n",
    "        output[fasta_id]=pred_with_idx\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('multi_modal_binding/results/inference/2024-09-11-18-21/predictions.pkl','rb') as f:\n",
    "    predictions=pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'sea': [(0.6507096, 249),\n",
      "         (0.48361024, 211),\n",
      "         (0.38016215, 138),\n",
      "         (0.24761724, 74),\n",
      "         (0.20594287, 85),\n",
      "         (0.15221491, 110),\n",
      "         (0.14734414, 30),\n",
      "         (0.14564994, 68),\n",
      "         (0.13472505, 156),\n",
      "         (0.13062719, 130)],\n",
      " 'see': [(0.690493, 249),\n",
      "         (0.5249078, 211),\n",
      "         (0.480403, 138),\n",
      "         (0.3035669, 188),\n",
      "         (0.26447755, 185),\n",
      "         (0.26083335, 251),\n",
      "         (0.22311726, 130),\n",
      "         (0.1829277, 85),\n",
      "         (0.17048396, 224),\n",
      "         (0.14448115, 110)]}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(best_residues('ZN','multi_modal_binding/results/inference/2024-09-11-18-21/predictions.pkl',10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "预测的结果认为sea和see 锌离子均与138/211/249号残基有最高的结合概率.\n",
    "\n",
    "实验测定结果中sea的25/211/249/251(加上未测定的前24个残基)号残基与Cd(替换Zn)离子直接配位.预测大致准确."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.生成3D结构\n",
    "\n",
    "首先将pkl预测文件按照阈值生成可能结合的残基序列:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: parse_dl_results.py [-h] [--lower_factor LOWER_FACTOR]\n",
      "                           {CA,ZN,MG,MN,FE,CU,CO} predictions_file\n",
      "\n",
      "Parse the deep learning predictions for real-world application.\n",
      "\n",
      "positional arguments:\n",
      "  {CA,ZN,MG,MN,FE,CU,CO}\n",
      "                        Type of ion to analyze.\n",
      "  predictions_file      Path to the predictions pickle file.\n",
      "\n",
      "options:\n",
      "  -h, --help            show this help message and exit\n",
      "  --lower_factor LOWER_FACTOR\n",
      "                        Lower factor to adjust the threshold. Set to 1.0 to\n",
      "                        use the original threshold.\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/3D_modeling/src/parse_dl_results.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "用法: parse_dl_results.py [-h] [--lower_factor LOWER_FACTOR]\n",
    "                           {CA,ZN,MG,MN,FE,CU,CO} predictions_file\n",
    "\n",
    "- 位置参数\n",
    "  - 离子类型:{CA,ZN,MG,MN,FE,CU,CO}\n",
    "  - predictions_file:pkl文件路径\n",
    "- 可选参数\n",
    "  - --help:显示帮助信息\n",
    "  - --lower_factor LOWER_FACTOR:下调阈值的缩放因子,默认1.0(即原阈值)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of parsed sequences is 2 for ZN\n",
      "Used threshold: 0.46\n",
      "Parsed results saved to: multi_modal_binding/results/inference/2024-09-11-18-21/parsed_result_predictions_ZN_lower_factor_0.60.pkl\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/3D_modeling/src/parse_dl_results.py ZN multi_modal_binding/results/inference/2024-09-11-18-21/predictions.pkl --lower_factor 0.6\n",
    "# 源代码ZN阈值为0.76,SEA/SEE前2位概率0.48左右,因此缩放到0.45"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后进行格点搜索/能量最小化等步骤,生成包含金属离子的3D结构."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: main.py [-h] [--debug] [--no-gpu] [--gpu_index GPU_INDEX]\n",
      "               [--pdb-dir PDB_DIR] [--output-dir OUTPUT_DIR]\n",
      "               [--prediction_result PREDICTION_RESULT] [--ion ION]\n",
      "               [--restraint_force_constant RESTRAINT_FORCE_CONSTANT]\n",
      "\n",
      "Process debug and GPU usage flags.\n",
      "\n",
      "options:\n",
      "  -h, --help            show this help message and exit\n",
      "  --debug               Enable debug mode.\n",
      "  --no-gpu              Disable GPU usage.\n",
      "  --gpu_index GPU_INDEX\n",
      "                        Specify the GPU ID to use. Default is 0.\n",
      "  --pdb-dir PDB_DIR     Directory to search for PDB files. Default is the\n",
      "                        current directory.\n",
      "  --output-dir OUTPUT_DIR\n",
      "                        Directory to save the results. Default is the current\n",
      "                        directory.\n",
      "  --prediction_result PREDICTION_RESULT\n",
      "                        the prediction results from deep learning model\n",
      "  --ion ION             Specify the ion type.\n",
      "  --restraint_force_constant RESTRAINT_FORCE_CONSTANT\n",
      "                        Specify how much restraint to put on the metal ion and\n",
      "                        backbone. Default is 83680 kJ/(mol nm^2).\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/3D_modeling/src/main.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "用法: main.py [-h] [--debug] [--no-gpu] [--gpu_index GPU_INDEX]\n",
    "               [--pdb-dir PDB_DIR] [--output-dir OUTPUT_DIR]\n",
    "               [--prediction_result PREDICTION_RESULT] [--ion ION]\n",
    "               [--restraint_force_constant RESTRAINT_FORCE_CONSTANT]\n",
    "\n",
    "选项:\n",
    "\n",
    "- --debug:显示调试信息\n",
    "- --no-gpu:禁用GPU使用\n",
    "- --gpu_index GPU_INDEX:指定使用的GPU设备编号,默认0\n",
    "- --pdb-dir PDB_DIR:搜索PDB文件的目录,默认当前目录\n",
    "- --output-dir PUTPUT_DIR:保存结果的目录,默认当前目录\n",
    "- --prediction_result PREDICTION_RESULT:解析过的预测文件路径\n",
    "- --ion ION:指定离子类型\n",
    "- --restraint_force_constant RESTRAINT_FORCE_CONSTANT:对金属离子和骨架的约束强度,默认83680 kJ/(mol nm^2).\n",
    "\n",
    "`src/main.py`脚本有小bug,建议用本笔记中的main.py替换,应该很快会修复."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processing sea\n",
      "Output file /home/regen/论文笔记/ESM_Bind_metal/sea_with_ZN_minimized.cif already exists. Skipping processing for sea.\n",
      "processing see\n",
      "Output file /home/regen/论文笔记/ESM_Bind_metal/see_with_ZN_minimized.cif already exists. Skipping processing for see.\n"
     ]
    }
   ],
   "source": [
    "!python $esm_bind_path/3D_modeling/src/main.py --ion ZN --restraint_force_constant 41840 --prediction_result multi_modal_binding/results/inference/2024-09-11-18-21/parsed_result_predictions_ZN_lower_factor_0.60.pkl --no-gpu\n",
    "# 41840为官方仓库中run-3d-modeling.sh使用的参数\n",
    "# 注意预测文件是解析过的pkl,不是inference.py生成的pkl\n",
    "# 由于笔记最开始禁用了gpu,此处不使用--no-gpu可能较慢,可以考虑取消禁用,因为是由openmm库进行调用\n",
    "# 使用CPU,生成2个序列结构耗时25min左右"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3 可视化\n",
    "\n",
    "使用pymol脚本,概率越高越红.由于本笔记环境不兼容pymol,仅生成pymol脚本,在其他环境运行pymol脚本."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "def py_for_esmbind_aa(probability: dict[str,list[tuple[float,int]]],fasta_id:str,pred_3d: str,reference_pdbid:str='',reference_pdb_chain:str='A',pred_ion:str='ZN',ref_ion:str='ZN') -> None:\n",
    "    '''根据esmbind预测的残基结合概率生成pymol脚本文件,对蛋白表面进行着色.在当前目录保存到`pred_3d.predicted.py`文件.\n",
    "\n",
    "    probability:{fasta_id:[(概率,残基序号)]}\n",
    "\n",
    "    fasta_id:序列的fasta id\n",
    "\n",
    "    pred_3d:esmbind预测的3D结构文件\n",
    "\n",
    "    reference_pdbid:参考的实验结构pdbid,如果未指定,则不作额外比较\n",
    "\n",
    "    reference_pdb_chain:参考的实验结构链号,默认为A链\n",
    "\n",
    "    pred_ion:预测结构文件中离子的残基名\n",
    "\n",
    "    ref_ion:参考结构文件中离子的残基名\n",
    "    '''\n",
    "\n",
    "    output_name=os.path.basename(pred_3d) +'.predicted.py'\n",
    "    reference_template='\\n'.join([\n",
    "        f'cmd.fetch(\"{reference_pdbid+reference_pdb_chain}\",\"ref\")',\n",
    "        'cmd.remove(\"solvent\")',\n",
    "        'cmd.color(\"greencyan\", \"ref\")',\n",
    "        f'cmd.color(\"green\", \"ref and resn {ref_ion}\")',\n",
    "        'cmd.align(\"pred\",\"ref\")',\n",
    "        f'cmd.distance(\"pred and resn {pred_ion}\",\"ref and resn {ref_ion}\")'\n",
    "        ])\n",
    "\n",
    "    color_template=[]\n",
    "    for p,res_id in probability[fasta_id]:\n",
    "        color_template.append(f'cmd.set_color(\"esmbind{res_id}\", {[255, 255-int(255 * p), 255-int(255 * p)]})')\n",
    "        color_template.append(f'''cmd.color(\"esmbind{res_id}\", \"pred and resi {res_id}\")''')\n",
    "    color_template='\\n'.join(color_template)\n",
    "\n",
    "    template = f'''import pymol.cmd as cmd\n",
    "cmd.load('{pred_3d}', 'pred')\n",
    "cmd.show('surface', 'pred')\n",
    "cmd.color('white', 'pred')\n",
    "cmd.set('transparency', 0.2, 'pred')\n",
    "{color_template}\n",
    "cmd.color('blue', 'pred and resn {pred_ion}')\n",
    "\n",
    "{reference_template}\n",
    "\n",
    "cmd.center()\n",
    "\n",
    "'''\n",
    "\n",
    "    with open(f'{output_name}.py','w+',encoding='utf-8') as outputfile:\n",
    "        outputfile.write(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_dir='multi_modal_binding/results/inference/2024-09-11-18-21'\n",
    "py_for_esmbind_aa(best_residues('ZN',pred_dir+'/predictions.pkl'),'sea','sea_with_ZN_minimized.cif','1esf',ref_ion='CD')\n",
    "# 1esf测定使用Cd2+替换Zn2+,方便结晶测定"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在当前目录打开终端,命令行使用`pymol xxxxxpredicted.py.py`或者在pymol图形化界面打开py文件即可.\n",
    "\n",
    "可以看到预测位置基本正确,预测离子(蓝)与实验测定位置(绿)相差1.8埃,可以接受.\n",
    "\n",
    "![](sea_pred.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r multi_modal_binding/\n",
    "!rm inference.json *.log *.npy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r esm/ esm_if/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm *_with_ZN_minimized.cif *.predicted.py.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
